import numpy as np
import opuslib
import zlib
import random

# 全局配置
SAMPLE_RATE = 24000
CHANNELS = 1
FRAME_SIZE = 120  # 每帧样本数（5ms 帧长）
SAMPLES_PER_FRAME_48K = 240  # 48kHz采样率下的每帧样本数（24000Hz时每帧120样本→48kHz对应240样本）

## 1. PCM处理工具
def float32_to_int16(float_array):
    float_array = np.clip(float_array, -1.0, 1.0)
    return (float_array * 32767).astype(np.int16)

## 2. Opus编码处理
def create_opus_encoder():
    """创建Opus编码器"""
    enc = opuslib.Encoder(SAMPLE_RATE, CHANNELS, opuslib.APPLICATION_AUDIO)
    enc.bitrate = 24000
    enc.complexity = 6
    enc.signal = opuslib.SIGNAL_VOICE
    return enc

def float32_pcm_to_opus(float32_pcm: np.ndarray, encoder=None) -> bytes:
    int16_data = float32_to_int16(float32_pcm)
    return int16_pcm_to_opus(int16_data, encoder)

def int16_pcm_to_opus(int16_pcm: np.ndarray, encoder=None) -> bytes:
    enc = encoder or create_opus_encoder()
    if len(int16_pcm.shape) > 1:
        int16_pcm = int16_pcm.flatten()
    return enc.encode(int16_pcm.tobytes(), FRAME_SIZE)

## 3. Ogg封装实现
def create_opus_header():
    """创建Opus头（ID Header和Comment Header）"""
    # ID Header结构：https://datatracker.ietf.org/doc/html/rfc7845#section-5.1
    id_header = b"OpusHead"  # 固定签名
    id_header += b"\x01"     # 版本 (1)
    id_header += bytes([CHANNELS])  # 声道数
    id_header += (384).to_bytes(2, "little")  # 预跳过(小端)
    id_header += SAMPLE_RATE.to_bytes(4, "little")  # 原始采样率
    id_header += (0).to_bytes(2, "little")    # 输出增益(小端)
    id_header += b"\x00"     # 声道映射
    
    # Comment Header：https://datatracker.ietf.org/doc/html/rfc7845#section-5.2
    vendor = b"ServerOpusStream"
    comment_header = b"OpusTags"                            # 固定签名
    comment_header += len(vendor).to_bytes(4, "little")    # 供应商字符串长度
    comment_header += vendor                                # 供应商字符串
    comment_header += (0).to_bytes(4, "little")            # 用户注释数量(0)
    
    return id_header, comment_header

def build_ogg_page(packet: bytes, page_index: int, granule_pos: int, stream_serial: int, bos=False, eos=False):
    """构建单个Ogg页面"""
    # 1. 计算段表 (segment table)
    packet_len = len(packet)
    segment_count = (packet_len + 255) // 256  # 每段最多255字节
    segment_table = []
    
    # 生成段长度表
    for i in range(segment_count):
        start = i * 255
        end = min(start + 255, packet_len)
        segment_table.append(end - start)
    
    # 2. 构建页面头部
    header_type = 0
    if bos: 
        header_type |= 1  # BOS标志
    if eos: 
        header_type |= 2  # EOS标志
        
    header = b"OggS"                  # 魔数
    header += b"\x00"                 # 版本
    header += bytes([header_type])    # 头部类型
    header += granule_pos.to_bytes(8, "little")      # 时间戳（48kHz采样率）
    header += stream_serial.to_bytes(4, "little")    # 流序列号
    header += page_index.to_bytes(4, "little")        # 页面序号
    header += b"\x00\x00\x00\x00"     # 校验和占位
    header += bytes([len(segment_table)])  # 段数量
    
    # 添加段表
    for seg_len in segment_table:
        header += bytes([seg_len])
    
    # 3. 构建完整页面（含数据）
    full_page = header + packet
    
    # 4. 计算并填充CRC32校验和
    checksum = zlib.crc32(full_page) & 0xFFFFFFFF
    full_page = header[:22] + checksum.to_bytes(4, "little") + header[26:] + packet
    
    return full_page

## 4. 流式处理封装
def generate_ogg_opus_stream(pcm_generator):
    """
    流式生成Ogg Opus格式音频
    :param pcm_generator: 生成器，每次yield一帧PCM数据 (float32, [-1,1]) 
    :yield: Ogg Opus格式数据流
    """
    # 创建唯一序列号
    stream_serial = random.randint(0, 0xFFFFFFFF)
    
    # 准备Opus标头
    id_header, comment_header = create_opus_header()
    
    # 构建初始头部页面
    yield build_ogg_page(
        packet=id_header,
        page_index=0,
        granule_pos=0,
        stream_serial=stream_serial,
        bos=True
    )
    
    yield build_ogg_page(
        packet=comment_header,
        page_index=1,
        granule_pos=0,
        stream_serial=stream_serial
    )
    
    # 创建Opus编码器
    encoder = create_opus_encoder()
    
    # 流式处理PCM帧
    page_index = 2
    granule_pos = 0
    
    for pcm_frame in pcm_generator:
        # 确保PCM帧长度正确
        if len(pcm_frame) != FRAME_SIZE * CHANNELS:
            raise ValueError(f"无效帧长，期望{FRAME_SIZE * CHANNELS}样本，实际{len(pcm_frame)}")
        
        # PCM转Opus
        opus_packet = float32_pcm_to_opus(pcm_frame, encoder)
        
        # 更新时间戳（48kHz下的采样位置）
        granule_pos += SAMPLES_PER_FRAME_48K
        
        # 构建Ogg页面
        yield build_ogg_page(
            packet=opus_packet,
            page_index=page_index,
            granule_pos=granule_pos,
            stream_serial=stream_serial
        )
        page_index += 1
    
    # 添加结束标记
    yield build_ogg_page(
        packet=b'',
        page_index=page_index,
        granule_pos=granule_pos,
        stream_serial=stream_serial,
        eos=True
    )